package com.zhaiker.keypointestimation

import android.content.Context
import android.graphics.Bitmap
import android.util.Log
import android.widget.Toast
import com.cv.tnn.model2.Detector
import com.cv.tnn.model2.FrameInfo
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import java.io.File
import java.io.FileOutputStream
import java.io.IOException

object KeypointEstimationHelper {

    private const val pose_model1: String = "pose/litehrnet18_17_192_256.sim"

    private const val pose_model2: String = "pose/mobilenet_v2_17_192_256.sim"

    private const val dets_model: String = "yolov5/yolov5s05_320.sim"

    private var TNN_MODEL_FILES: Array<String> = arrayOf(
        pose_model1,
        pose_model2,
        dets_model,
    )

    private var USE_GPU: Boolean = false

    private var MODEL_ID: Int = 0

    private var NUM_THREAD: Int = 1

    private var _sdkAvailable : Boolean = false

    //sdk是否可用
    val sdkAvailable get() = _sdkAvailable

    private const val TAG = "keypointEstimation"

    private suspend fun copyModelFromAssetsToData(context: Context) : Boolean{
        // assets目录下的模型文件名
        //Toast.makeText(context, "Copy model to data...", Toast.LENGTH_SHORT).show()
        try {
            for (tnn_model in TNN_MODEL_FILES) {
                //  {"face/rfb1.0_face_320_320.opt.tnnproto", "face/rfb1.0_face_320_320.opt.tnnmodel"};
                val tnnproto = "$tnn_model.tnnproto"
                val tnnmodel = "$tnn_model.tnnmodel"
                Log.w(TAG, "copy model:$tnn_model")
                copyAssetFileToFiles(context, tnnproto)
                copyAssetFileToFiles(context, tnnmodel)
            }

            //Toast.makeText(context, "Copy model Success", Toast.LENGTH_SHORT).show()
            return true
        } catch (e: Exception) {
            e.printStackTrace()
            Toast.makeText(context, "Copy keypoint model Error", Toast.LENGTH_SHORT).show()
            return false
        }
    }

    @Throws(IOException::class)
    private suspend fun copyAssetFileToFiles(context: Context, filename: String) = withContext(Dispatchers.IO){
        val file = File(context.filesDir.toString() + File.separator + filename)
        if(file.exists()){
            Log.i("object-detection-tnn","${filename} exists, no copy required")
            return@withContext
        }

        val parent = File(file.parent)
        if (!parent.exists()) {
            parent.mkdirs()
        }
        file.createNewFile()

        val inputStream = context.assets.open(filename)
        val buffer = ByteArray(inputStream.available())
        inputStream.read(buffer)
        inputStream.close()


        val os = FileOutputStream(file)
        os.write(buffer)
        os.close()
    }

    private fun initModel(context: Context,) {
        val root: String = context.filesDir.toString() + File.separator
        val length: Int = TNN_MODEL_FILES.size
        val pose_model: String =
            TNN_MODEL_FILES[MODEL_ID] // 是pose模型
        val dets_model: String = TNN_MODEL_FILES[length - 1] // 默认最后一个检测模型
        Detector.init(
            dets_model,
            pose_model,
            root,
            MODEL_ID,
            NUM_THREAD,
            USE_GPU
        )
    }

    /**
     * @param modelId 0:litehrnet,速度快;1:mobilenet_v2,精度高
     */
    suspend fun init(context: Context,modelId : Int){
        if(!copyModelFromAssetsToData(context)){
            return
        }
        MODEL_ID = modelId
        initModel(context)
        _sdkAvailable = true
    }

    /**
     * @param bitmap 图像（bitmap），ARGB_8888格式
     * @param score_threshold 置信度阈值
     * @param iou_threshold IOU阈值
     * @param pose_threshold 关键点阈值
     * @return 返回一个数组，如果没有检测到人体，会返回一个size为0的数组，不会返回null
     */
    fun detect(bitmap : Bitmap,score_threshold : Float = 0.5f,iou_threshold : Float = 0.5f,pose_threshold : Float = 0.3f) : Array<FrameInfo>{
        return Detector.detect(bitmap,score_threshold, iou_threshold, pose_threshold)
    }

    /**
     * 检测最大的人体
     */
    fun detectFirst(bitmap : Bitmap,score_threshold : Float = 0.5f,iou_threshold : Float = 0.5f,pose_threshold : Float = 0.3f) : FrameInfo?{
        val result = Detector.detect(bitmap,score_threshold, iou_threshold, pose_threshold)
        if(result.isNullOrEmpty()){
            return null
        }
        return result[0]

    }

}